#pragma once

#include <stx/core/algo.hpp>

inline namespace stx
{
template<IsAnyType T = void>
using CoroHandle = std::coroutine_handle<T>;

template<IsAnyType T>
class SharedFuture;

template<IsAnyType T>
class Future;

template<IsAnyType>
struct IsSharedFuture__ final : std::false_type
{};

template<IsAnyType T>
struct IsSharedFuture__<SharedFuture<T>> final : std::true_type
{};

template<typename T>
concept IsSharedFuture = IsSharedFuture__<NoCvr<T>>::value;

template<IsAnyType>
struct IsFuture__ final : std::false_type
{};

template<IsAnyType T>
struct IsFuture__<Future<T>> final : std::true_type
{};

template<typename T>
concept IsFuture = IsFuture__<NoCvr<T>>::value;

template<typename T>
concept IsStxFuture = AnyTrue<IsSharedFuture<T>, IsFuture<T>>;

template<IsAnyType>
struct IsStdSharedFuture__ final : std::false_type
{};

template<IsAnyType T>
struct IsStdSharedFuture__<std::shared_future<T>> final : std::true_type
{};

template<IsAnyType>
struct IsStdFuture__ final : std::false_type
{};

template<IsAnyType T>
struct IsStdFuture__<std::future<T>> final : std::true_type
{};

template<typename T>
concept IsStdFuture =
    AnyTrue<IsStdFuture__<T>::value, IsStdSharedFuture__<T>::value>;

template<IsAnyType T>
struct StxFutureStripper__ final
{
    using type = T;
};

template<IsAnyType T>
requires EXPR(!IsStxFuture<T>)
struct StxFutureStripper__<Future<T>> final
{
    using type = T;
};

template<IsAnyType T>
requires EXPR(!IsStxFuture<T>)
struct StxFutureStripper__<SharedFuture<T>> final
{
    using type = T;
};

template<IsAnyType T>
using StripStxFuture = typename StxFutureStripper__<T>::type;

using Continuations__         = std::deque<Task>;
using AtomicContinuations__   = Atomic<Continuations__>;
using SpAtomicContinuations__ = std::shared_ptr<AtomicContinuations__>;

template<IsAnyType T = void>
class Promise final
{
    std::promise<T>         prom_;
    SpAtomicContinuations__ sp_atom_conts_ =
        std::make_shared<AtomicContinuations__>();

public:
    using ValueType_ = T;

    Promise()
    {
        this->init_();
    }

    template<IsNonVoid Allocator_>
    Promise(std::allocator_arg_t, Allocator_ const& alloc)
        : prom_(std::allocator_arg, alloc)
    {
        this->init_();
    }

    NON_COPYABLE(Promise);

    DEFINE_MOVE_SEMANTICS(Promise, {
        this->prom_          = std::exchange(other.prom_, {});
        this->sp_atom_conts_ = std::exchange(other.sp_atom_conts_, {});
    });

    [[nodiscard]] bool isReady() const
    {
        return this->sp_atom_conts_->visit([](auto const& conts)
                                           { return conts.empty(); });
    }

    [[nodiscard]] Future<T> getFuture();

    void setException(char const* const sz_err_msg)
    {
        try
        {
            throw std::runtime_error(sz_err_msg);
        }
        catch (...)
        {
            this->setException(std::current_exception());
        }
    }

    void setException(std::exception_ptr const eptr)
    {
        this->prom_.set_exception(eptr);
        this->callConts_();
    }

    void setValue(IsNonVoid auto&&... value)
    {
        if constexpr (sizeof...(value) == 0)
        {
            this->prom_.set_value();
        }
        else if constexpr (sizeof...(value) == 1)
        {
            static_assert(IsNonVoid<T>);
            static_assert(std::constructible_from<T, decltype(value)...>);

            this->prom_.set_value(FORWARD(value)...);
        }
        else
        {
            static_assert(AlwaysFalse<T>);
        }

        this->callConts_();
    }

private:
    void init_()
    {
        this->sp_atom_conts_->visit([](auto& conts) { insert(conts, [] {}); });
    }

    void callConts_()
    {
        auto tmp_conts = this->sp_atom_conts_->visit(
            [](auto& conts) { return std::exchange(conts, {}); });
        ASSERT(!tmp_conts.empty());

        ranges::for_each(tmp_conts, BIND(std::invoke, _1));
    }
};

template<IsStxFuture StxFuture_, IsAnyType T>
class CommonCoroPromise
{
protected:
    Promise<T> prom_;

public:
    [[nodiscard]] static auto get_return_object_on_allocation_failure()
    {
        return StxFuture_{};
    }

    [[nodiscard]] auto get_return_object()
    {
        if constexpr (IsFuture<StxFuture_>)
        {
            return this->prom_.getFuture();
        }
        else
        {
            return this->prom_.getFuture().share();
        }
    }

    [[nodiscard]] std::suspend_never initial_suspend() noexcept
    {
        return {};
    }

    [[nodiscard]] std::suspend_never final_suspend() noexcept
    {
        return {};
    }

    void unhandled_exception()
    {
        this->prom_.setException(std::current_exception());
    }
};

template<IsStxFuture StxFuture_, IsAnyType T>
class CoroPromise final : public CommonCoroPromise<StxFuture_, T>
{
public:
    template<IsNonVoid U>
    requires EXPR(std::constructible_from<T, U&&>)
    void return_value(U&& value)
    {
        this->prom_.setValue(FORWARD(value));
    }
};

template<IsStxFuture StxFuture_>
class CoroPromise<StxFuture_, void> final
    : public CommonCoroPromise<StxFuture_, void>
{
public:
    void return_void()
    {
        this->prom_.setValue();
    }
};

template<IsStdFuture StdFuture, IsAnyType T = void>
requires EXPR(!IsStxFuture<T>)
class CommonFuture
{
protected:
    StdFuture               std_fut_;
    SpAtomicContinuations__ sp_atom_conts_;

public:
    [[nodiscard]] bool await_ready() const noexcept
    {
        return this->isReady();
    }

    void await_suspend(auto h_coro) const noexcept
    {
        this->registerContinuation_([h_coro] mutable { h_coro.resume(); });
    }

    [[nodiscard]] bool isValid() const noexcept
    {
        return this->std_fut_.valid();
    }

    [[nodiscard]] bool isReady() const noexcept
    {
        return std::future_status::ready == this->waitFor(0ms);
    }

    void wait() const
    {
        this->std_fut_.wait();
    }

    // NOLINTNEXTLINE
    std::future_status waitFor(IsStdDuration auto const dur) const
    {
        return this->std_fut_.wait_for(dur);
    }

    // NOLINTNEXTLINE
    std::future_status waitUntil(IsStdTimePoint auto const tp) const
    {
        return this->std_fut_.wait_until(tp);
    }

protected:
    CommonFuture() noexcept = default;

    CommonFuture(StdFuture std_fut, SpAtomicContinuations__ sp_atom_conts)
        : std_fut_(std::move(std_fut))
        , sp_atom_conts_(std::move(sp_atom_conts))
    {
        ASSERT(this->std_fut_.valid());
        ASSERT(this->sp_atom_conts_);
    }

    void registerContinuation_(Task const cb_cont) const
    {
        ASSERT(this->sp_atom_conts_);

        auto ready = false;
        this->sp_atom_conts_->visit(
            [&ready, cb_cont](auto& cb_conts)
            {
                if (std::empty(cb_conts))
                {
                    ready = true;
                }
                else
                {
                    insert(cb_conts, cb_cont);
                }
            });

        if (ready)
        {
            cb_cont();
        }
    }

    template<IsNonVoid Fn_, IsAnyType Arg_>
    struct ReturnType_ final
    {
        using type = std::invoke_result_t<Fn_, Arg_&&>;
    };

    template<IsNonVoid Fn_, IsNonVoid Arg_>
    struct ReturnType_<Fn_, Arg_&> final
    {
        using type = std::invoke_result_t<Fn_, Arg_&>;
    };

    template<IsNonVoid Fn_>
    struct ReturnType_<Fn_, void> final
    {
        using type = std::invoke_result_t<Fn_>;
    };

#define SET_VALUE__(...)                                 \
    try                                                  \
    {                                                    \
        sp_prom->setValue(__VA_ARGS__);                  \
    }                                                    \
    catch (...)                                          \
    {                                                    \
        sp_prom->setException(std::current_exception()); \
    }                                                    \
    static_assert(true)

    template<IsStxFuture RealFutType_>
    auto then_(IsNonVoid auto const cb_cont, IsNonVoid auto const sp_self)
    {
        ASSERT(this->isValid());
        ASSERT(this->sp_atom_conts_);

        using R0 = std::conditional_t<
            IsFuture<RealFutType_>,
            decltype(std::invoke(DECAY_COPY(cb_cont), std::move(*sp_self))),
            decltype(std::invoke(DECAY_COPY(cb_cont), *sp_self))>;
        using R1 = NoCvr<R0>;
        using R2 = StripStxFuture<R1>;

        auto const sp_prom = std::make_shared<Promise<R2>>();

        auto const fn_cont_internal =
            [DECAY_COPY_ARG(cb_cont), sp_self, sp_prom]
        {
            auto const fn_get_self_obj = [sp_self] -> decltype(auto)
            {
                if constexpr (IsFuture<RealFutType_>)
                {
                    return std::move(*sp_self);
                }
                else
                {
                    static_assert(IsSharedFuture<RealFutType_>);
                    return *sp_self;
                }
            };

            if constexpr (IsStxFuture<R0>)
            {
                if constexpr (IsNonVoid<R2>)
                {
                    std::invoke(cb_cont, fn_get_self_obj())
                        .then([sp_prom](auto&& v) { SET_VALUE__(v.get()); });
                }
                else
                {
                    std::invoke(cb_cont, fn_get_self_obj())
                        .then([sp_prom](auto&&) { SET_VALUE__(); });
                }
            }
            else
            {
                if constexpr (IsNonVoid<R2>)
                {
                    SET_VALUE__(std::invoke(cb_cont, fn_get_self_obj()));
                }
                else
                {
                    std::invoke(cb_cont, fn_get_self_obj());
                    SET_VALUE__();
                }
            }
        };

        sp_self->registerContinuation_(fn_cont_internal);
        return sp_prom->getFuture();
    }

    template<IsStxFuture RealFutType_>
    auto then_(IsNonVoid auto const cb_cont) &
    {
        return this->template then_<RealFutType_>(
            cb_cont, static_cast<RealFutType_*>(this));
    }

    template<IsStxFuture RealFutType_>
    auto then_(IsNonVoid auto const cb_cont) &&
    {
        auto const p_self  = static_cast<RealFutType_*>(this);
        auto       sp_self = std::shared_ptr<RealFutType_>{};
        if constexpr (IsFuture<RealFutType_>)
        {
            sp_self = std::make_shared<RealFutType_>(std::move(*p_self));
        }
        else
        {
            sp_self = std::make_shared<RealFutType_>(*p_self);
        }

        return this->template then_<RealFutType_>(cb_cont, sp_self);
    }
};

template<IsAnyType T = void>
class SharedFuture final : public CommonFuture<std::shared_future<T>, T>
{
    friend Future<T>;
    using Base_ = CommonFuture<std::shared_future<T>, T>;

    SharedFuture(std::shared_future<T>         sfut,
                 SpAtomicContinuations__ const sp_atom_conts)
        : Base_(std::move(sfut), sp_atom_conts)
    {}

public:
    using promise_type = CoroPromise<SharedFuture, T>;
    using ValueType_   = T;

    SharedFuture() noexcept = default;

    decltype(auto) get() const noexcept // NOLINT
    {
        return this->std_fut_.get();
    }

    decltype(auto) await_resume() const noexcept // NOLINT
    {
        return this->get();
    }

    auto then(IsNonVoid auto const cb_cont) &
    {
        return this->template then_<SharedFuture>(cb_cont);
    }

    auto then(IsNonVoid auto const cb_cont) &&
    {
        return this->template then_<SharedFuture>(cb_cont);
    }
};

template<IsAnyType T = void>
class Future final : public CommonFuture<std::future<T>, T>
{
    friend Promise<T>;
    using Base_ = CommonFuture<std::future<T>, T>;

    Future(std::future<T> fut, SpAtomicContinuations__ sp_atom_conts)
        : Base_(std::move(fut), sp_atom_conts)
    {}

public:
    using promise_type = CoroPromise<Future, T>;
    using ValueType_   = T;

    Future() noexcept = default;

    NON_COPYABLE(Future);

    DEFINE_MOVE_SEMANTICS(Future, {
        this->std_fut_       = std::move(other.std_fut_);
        this->sp_atom_conts_ = std::exchange(other.sp_atom_conts_, {});
    });

    decltype(auto) get() noexcept
    {
        return this->std_fut_.get();
    }

    decltype(auto) await_resume() noexcept
    {
        return this->get();
    }

    [[nodiscard]] SharedFuture<T> share() noexcept
    {
        return SharedFuture<T>(this->std_fut_.share(), this->sp_atom_conts_);
    }

    auto then(IsNonVoid auto const cb_cont) &
    {
        return this->template then_<Future>(cb_cont);
    }

    auto then(IsNonVoid auto const cb_cont) &&
    {
        return this->template then_<Future>(cb_cont);
    }
};

template<IsAnyType T>
[[nodiscard]] Future<T> Promise<T>::getFuture()
{
    return Future<T>(this->prom_.get_future(), this->sp_atom_conts_);
}

inline decltype(auto) operator|(IsStxFuture auto fut, IsNonVoid auto const fn)
{
    return fut.then(fn);
}

[[nodiscard]] inline auto makeReadyFuture(IsNonVoid auto&& value)
{
    auto prom = Promise<NoCvr<decltype(value)>>{};
    prom.setValue(FORWARD(value));
    return prom.getFuture();
}

[[nodiscard]] inline auto makeReadyFuture()
{
    auto prom = Promise<>{};
    prom.setValue();
    return prom.getFuture();
}

inline decltype(auto) forwardStxFuture__(IsStxFuture auto&& v)
{
    if constexpr (IsSharedFuture<decltype(v)>)
    {
        auto const& r = v;
        return r;
    }
    else
    {
        return std::move(v); // NOLINT
    }
};

struct WhenAllNotifier__ final
{
    Promise<>   prom;
    std::size_t pending_fut_cnt = {};
};

template<std::input_iterator InputIter_,
         IsNonVoid           Futures_ =
             std::vector<typename std::iterator_traits<InputIter_>::value_type>>
requires requires(InputIter_ pos)
{
    requires IsStxFuture<decltype(*pos)>;
}
[[nodiscard]] Future<Futures_> whenAll(InputIter_ const first,
                                       InputIter_ const last)
{
    auto futs = Futures_{};
    for (auto pos = first; pos != last; ++pos)
    {
        if (IsSharedFuture<decltype(*pos)>)
        {
            insert(futs, *pos);
        }
        else
        {
            insert(futs, std::move(*pos));
        }
    }

    if (futs.empty())
    {
        co_return futs;
    }

    auto const sp_atomic_notifier =
        std::make_shared<Atomic<WhenAllNotifier__>>();
    sp_atomic_notifier->visit([&](auto& notifier)
                              { notifier.pending_fut_cnt = futs.size(); });

    for (auto& fut : futs)
    {
        fut.then(
            [sp_atomic_notifier](auto&&)
            {
                sp_atomic_notifier->visit(
                    [](auto& notifier)
                    {
                        if (0 == --notifier.pending_fut_cnt)
                        {
                            if (!notifier.prom.isReady())
                            {
                                notifier.prom.setValue();
                            }
                        }
                    });
            });
    }

    co_await sp_atomic_notifier->visit([](auto& notifier)
                                       { return notifier.prom.getFuture(); });
    co_return futs;
}

template<IsStxFuture... Futures_,
         IsNonVoid Tuple_ = std::tuple<NoCvr<Futures_>...>>
[[nodiscard]] Future<Tuple_> whenAll(Futures_&&... futs)
{
    if constexpr (sizeof...(futs) == 0)
    {
        co_return std::tuple<>{};
    }
    else
    {
        auto const sp_atomic_notifier =
            std::make_shared<Atomic<WhenAllNotifier__>>();
        auto fut_notifier = sp_atomic_notifier->visit(
            [](auto& notifier) { return notifier.prom.getFuture(); });
        sp_atomic_notifier->visit(
            [&](auto& notifier)
            { notifier.pending_fut_cnt = sizeof...(futs); });

        auto const fn_set_notifier = [=](auto& fut)
        {
            fut.then(
                [sp_atomic_notifier](auto&&)
                {
                    sp_atomic_notifier->visit(
                        [](auto& notifier)
                        {
                            if (0 == --notifier.pending_fut_cnt)
                            {
                                if (!notifier.prom.isReady())
                                {
                                    notifier.prom.setValue();
                                }
                            }
                        });
                });
        };

        auto tpl = std::make_tuple(forwardStxFuture__(futs)...);
        std::apply([&](auto&... fut) { (..., fn_set_notifier(fut)); }, tpl);

        co_await fut_notifier;
        co_return tpl;
    }
}

template<IsNonVoid Futures_>
requires IsStdTuple<Futures_> || IsStdVector<Futures_>
struct WhenAnyResult final
{
    std::ssize_t idx;
    Futures_     futs;
};

template<std::input_iterator InputIter_,
         IsNonVoid           Futures_ =
             std::vector<typename std::iterator_traits<InputIter_>::value_type>>
requires requires(InputIter_ pos)
{
    requires IsStxFuture<decltype(*pos)>;
}
[[nodiscard]] Future<WhenAnyResult<Futures_>> whenAny(InputIter_ const first,
                                                      InputIter_ const last)
{
    auto futs = Futures_{};
    for (auto pos = first; pos != last; ++pos)
    {
        if constexpr (IsSharedFuture<decltype(*pos)>)
        {
            insert(futs, *pos);
        }
        else
        {
            insert(futs, std::move(*pos));
        }
    }

    auto ret = WhenAnyResult<Futures_>{};
    if (futs.empty())
    {
        ret.idx = -1;
        co_return ret;
    }

    auto sp_atomic_prom = std::make_shared<Atomic<Promise<std::ssize_t>>>();
    auto fut_notifier =
        sp_atomic_prom->visit([](auto& prom) { return prom->getFuture(); });

    auto i = 0z;
    for (auto& fut : futs)
    {
        fut.then(
            [i, sp_atomic_prom](auto&&)
            {
                sp_atomic_prom->visit(
                    [i](auto& prom)
                    {
                        if (!prom.isReady())
                        {
                            prom.setValue(i);
                        }
                    });
            });

        ++i;
    }

    ret.idx  = co_await fut_notifier;
    ret.futs = std::move(futs);
    co_return ret;
}

template<IsStxFuture... Futures_,
         IsNonVoid Tuple_ = std::tuple<NoCvr<Futures_>...>>
[[nodiscard]] Future<WhenAnyResult<Tuple_>> whenAny(Futures_&&... futs)
{
    if constexpr (sizeof...(futs) == 0)
    {
        co_return WhenAnyResult<Tuple_>{.idx = -1, .futs = std::tuple<>{}};
    }
    else
    {
        auto tpl = std::make_tuple(forwardStxFuture__(futs)...);

        auto sp_atomic_prom_notifier =
            std::make_shared<Atomic<Promise<std::ssize_t>>>();
        auto fut_notifier = sp_atomic_prom_notifier->visit(
            [](auto& prom_notifier) { return prom_notifier.getFuture(); });

        auto       i           = 0z;
        auto const fn_set_cont = [&](auto& fut)
        {
            fut.then(
                [i, sp_atomic_prom_notifier](auto&&)
                {
                    sp_atomic_prom_notifier->visit(
                        [i](auto& prom_notifier)
                        {
                            if (!prom_notifier.isReady())
                            {
                                prom_notifier.setValue(i);
                            }
                        });
                });

            ++i;
        };

        std::apply([&](auto&... fut) { (..., fn_set_cont(fut)); }, tpl);

        auto ret = WhenAnyResult<Tuple_>{};
        ret.idx  = co_await fut_notifier;
        ret.futs = std::move(tpl);
        co_return ret;
    }
}

template<IsRawTransferable T>
class Generator final
{
    using FnGen_ = std::function<Future<std::optional<T>>()>;

    std::optional<T> opt_value_;
    FnGen_           fn_gen_;

public:
    class iterator final
    {
        Generator* p_owner_ = {};

    public:
        EQUALITY_OPERATOR(iterator);

        iterator() noexcept = default;

        iterator(Generator* const p_owner) : p_owner_(p_owner)
        {}

        [[nodiscard]] T& operator*() const
        {
            return *this->p_owner_->opt_value_;
        }

        [[nodiscard]] T* operator->() const
        {
            return &*this->p_owner_->opt_value_;
        }

        [[nodiscard]] Future<iterator> operator++()
        {
            if (this->p_owner_)
            {
                if (!co_await this->p_owner_->getNextValue_())
                {
                    this->p_owner_ = {};
                }
            }

            co_return *this;
        }

        void operator++(int) = delete;
    };

public:
    explicit Generator(FnGen_ fn_gen) : fn_gen_(std::move(fn_gen))
    {}

    [[nodiscard]] Future<iterator> begin()
    {
        if (co_await this->getNextValue_())
        {
            co_return iterator(this);
        }

        co_return iterator{};
    }

    [[nodiscard]] constexpr iterator end() const noexcept
    {
        return {};
    }

private:
    [[nodiscard]] Future<bool> getNextValue_()
    {
        if (auto opt_value = co_await this->fn_gen_())
        {
            this->opt_value_ = std::move(opt_value);
            co_return true;
        }

        this->opt_value_ = {};
        co_return false;
    }
};
} // namespace stx

/*
WARNING
-------
Don't use LOCK(...) to prom.setValue()!

auto g_prom_mtx = std::recursive_mutex{};

void neverDoLikeThis(Promise<>& prom)
{
    LOCK(g_prom_mtx);
    if (!prom.isReady())
    {
        prom.setValue();
    }
}

REASON
------
prom.setValue() will trigger many continuations, the lock's behavior is totally
uncontrolled; even the lock type is of std::recursive_mutex.

Future<> demoGenerator()
{
    auto fn_gen = [n = 0] mutable -> Future<std::optional<int>>
    {
        while (n < 10)
        {
            co_await 100ms;
            co_return n++;
        }

        co_return std::nullopt;
    };

    CORO_FOR (auto const n : Generator<int>(fn_gen))
    {
        xout(getOsTid(), _I_, n);
    }
}
*/